{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Segmentation with VGGNet U-Net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on code in Divam Gupta's [image-segmentation-keras](https://github.com/divamgupta/image-segmentation-keras) repository. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "! pip install opencv-contrib-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "! pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "! pip install keras_segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import keras_segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = keras_segmentation.models.unet.vgg_unet(n_classes=51, \n",
    "                                                      input_height=416, \n",
    "                                                      input_width=608)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 12/367 [00:00<00:02, 119.63it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Verifying train dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 367/367 [00:01<00:00, 203.14it/s]\n",
      " 21%|██        | 21/101 [00:00<00:00, 202.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset verified! \n",
      "Verifying val dataset\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 101/101 [00:00<00:00, 199.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset verified! \n",
      "Starting Epoch  0\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 139s - loss: 1.2426 - acc: 0.6802 - val_loss: 1.2152 - val_acc: 0.6010\n",
      "saved  model_output/image-seg.model.0\n",
      "Finished Epoch 0\n",
      "Starting Epoch  1\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.7111 - acc: 0.7843 - val_loss: 0.8228 - val_acc: 0.7302\n",
      "saved  model_output/image-seg.model.1\n",
      "Finished Epoch 1\n",
      "Starting Epoch  2\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.5134 - acc: 0.8424 - val_loss: 0.4538 - val_acc: 0.8648\n",
      "saved  model_output/image-seg.model.2\n",
      "Finished Epoch 2\n",
      "Starting Epoch  3\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.4178 - acc: 0.8713 - val_loss: 0.4496 - val_acc: 0.8576\n",
      "saved  model_output/image-seg.model.3\n",
      "Finished Epoch 3\n",
      "Starting Epoch  4\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 127s - loss: 0.3410 - acc: 0.8936 - val_loss: 0.3900 - val_acc: 0.8749\n",
      "saved  model_output/image-seg.model.4\n",
      "Finished Epoch 4\n",
      "Starting Epoch  5\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.2911 - acc: 0.9083 - val_loss: 0.3558 - val_acc: 0.8881\n",
      "saved  model_output/image-seg.model.5\n",
      "Finished Epoch 5\n",
      "Starting Epoch  6\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.2539 - acc: 0.9193 - val_loss: 0.3884 - val_acc: 0.8793\n",
      "saved  model_output/image-seg.model.6\n",
      "Finished Epoch 6\n",
      "Starting Epoch  7\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.2179 - acc: 0.9304 - val_loss: 0.3390 - val_acc: 0.8997\n",
      "saved  model_output/image-seg.model.7\n",
      "Finished Epoch 7\n",
      "Starting Epoch  8\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.1837 - acc: 0.9406 - val_loss: 0.2975 - val_acc: 0.9111\n",
      "saved  model_output/image-seg.model.8\n",
      "Finished Epoch 8\n",
      "Starting Epoch  9\n",
      "Epoch 1/1\n",
      "512/512 [==============================] - 128s - loss: 0.1614 - acc: 0.9472 - val_loss: 0.4474 - val_acc: 0.8777\n",
      "saved  model_output/image-seg.model.9\n",
      "Finished Epoch 9\n"
     ]
    }
   ],
   "source": [
    "model.train(\n",
    "    train_images = \"image-seg-data/images_prepped_train/\",\n",
    "    train_annotations = \"image-seg-data/annotations_prepped_train/\",\n",
    "    checkpoints_path = \"model_output/image-seg/\",\n",
    "    epochs = 10,\n",
    "    validate=True,\n",
    "    val_images = \"image-seg-data/images_prepped_test/\",\n",
    "    val_annotations = \"image-seg-data/annotations_prepped_test/\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "out = model.predict_segmentation(\n",
    "    inp=\"image-seg-data/images_prepped_test/0016E5_07965.png\",\n",
    "    out_fname=\"output-10epochs.png\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model.load_weights('model_output/image-seg.8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "out = model.predict_segmentation(\n",
    "    inp=\"image-seg-data/images_prepped_test/0016E5_07965.png\",\n",
    "    out_fname=\"output-9epochs.png\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
